import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms

from config import get_parser

'''
注意， 数据比那换的时候，千万不要pil 之后在np.array， 效率会变得非常慢
'''

f = open(get_parser().alphabet, 'r', encoding="utf-8")

alphabet = ''.join([s.strip('\n') for s in f.readlines()])

alphabet = '-' + alphabet


def convert_plate(target):
    plate = ''
    for i in target:
        plate += alphabet[i]
    return plate


data_transform = transforms.Compose([
    transforms.Resize([32, 140]),

    transforms.ToTensor(),

    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


class RecTextLineDataset(Dataset):
    def __init__(self, lines, type='train'):
        self.args = get_parser()
        self.alphabet = alphabet
        self.str2idx = {c: i for i, c in enumerate(self.alphabet)}
        self.labels = lines

        self.type = type

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, index):
        target = self.labels[index].split('/')[-1].strip('.jpg')
        img_path = self.labels[index]

        img = Image.open(img_path).convert("RGB")

        label = list()
        target = target.strip('\n')

        for c in target:
            if c == '\n':
                continue
            label.append(self.str2idx[c])

        return img, label, len(target), target


def recCollate(batch):
    imgs = []
    labels = []
    lengths = []
    labels_words = []
    for _, sample in enumerate(batch):
        img, label, length, words = sample
        img = data_transform(img)
        imgs.append(img)
        labels.extend(label)
        lengths.append(length)
        labels_words.append(words)
    labels = np.array(labels).flatten().astype(np.int)

    return (torch.stack(imgs, 0),
            torch.from_numpy(labels),
            lengths,
            labels_words)
